package com.example.springboot_demo.tools;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.alibaba.ttl.threadpool.TtlExecutors;

import java.util.Iterator;
import java.util.WeakHashMap;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

/**
 * @author
 * @since 2021/9/14 3:13 下午
 */
public class EsThreadLocal<T> extends InheritableThreadLocal<T> {

    private static InheritableThreadLocal<WeakHashMap<EsThreadLocal<Object>, ?>> holder = new InheritableThreadLocal<WeakHashMap<EsThreadLocal<Object>, ?>>() {
        @Override
        protected WeakHashMap<EsThreadLocal<Object>, ?> childValue(WeakHashMap<EsThreadLocal<Object>, ?> parentValue) {
            return new WeakHashMap<>(parentValue);
        }

        @Override
        protected WeakHashMap<EsThreadLocal<Object>, ?> initialValue() {
            return new WeakHashMap<>();
        }
    };

    @Override
    public int hashCode() {
        return super.hashCode();
    }

    @Override
    public boolean equals(Object obj) {
        return super.equals(obj);
    }

    @Override
    public T get() {
        T value = super.get();
        if (null != value)
            addToHolder();
        return value;
    }


    @Override
    public void set(T value) {
        if (null == value) {
            removeFromHolder();
            super.remove();
        } else {
            super.set(value);
            addToHolder();
        }
    }

    private void removeFromHolder() {

        holder.get().remove(this);
    }


    private void addToHolder() {

        if (!holder.get().containsKey(this))
            holder.get().put((EsThreadLocal<Object>) this, null);
    }

    static class SnapShot {
        final WeakHashMap<EsThreadLocal<Object>, Object> ctlValue;

        private SnapShot(WeakHashMap<EsThreadLocal<Object>, Object> ctlValue) {
            this.ctlValue = ctlValue;
        }
    }

    static class Transmitter {

        public static SnapShot capture() {
            return new SnapShot(captureCtlValues());
        }

        private static WeakHashMap<EsThreadLocal<Object>, Object> captureCtlValues() {
            return holder.get().keySet().stream().collect(Collectors.toMap(ctlItem -> ctlItem, EsThreadLocal::get, (a, b) -> b, WeakHashMap::new));
        }

        public static SnapShot replay(SnapShot snapShot) {
            WeakHashMap<EsThreadLocal<Object>, Object> capture = snapShot.ctlValue;
            WeakHashMap<EsThreadLocal<Object>, Object> backValue = new WeakHashMap<>();
            /*
             * 从holder中获取当前线程持有的threadLocal的Map,进行迭代保存
             */
            Iterator<EsThreadLocal<Object>> iterator = holder.get().keySet().iterator();
            while (iterator.hasNext()) {
                EsThreadLocal<Object> threadLocal = iterator.next();
                backValue.put(threadLocal, threadLocal.get());
                if (!capture.containsKey(threadLocal)) {
                    iterator.remove();
                    threadLocal.remove();
                }
            }
            /*
             设置上capture
             */
            setThreadLocal(capture);
            return new SnapShot(backValue);
        }


        public static void setThreadLocal(WeakHashMap<EsThreadLocal<Object>, Object> ctlValues) {
            ctlValues.forEach(EsThreadLocal::set);
        }

        public static void restore(EsThreadLocal.SnapShot backUp) {

            Iterator<EsThreadLocal<Object>> iterator = holder.get().keySet().iterator();

            while (iterator.hasNext()) {
                EsThreadLocal<Object> threadLocal = iterator.next();
                if (!backUp.ctlValue.containsKey(threadLocal)) {
                    iterator.remove();
                    threadLocal.remove();
                }
            }
            setThreadLocal(backUp.ctlValue);
        }
    }

    static class EsRunnable implements Runnable {

        private AtomicReference<SnapShot> captureRef;

        private Runnable runnable;

        public EsRunnable(Runnable runnable) {
            this.runnable = runnable;
            captureRef = new AtomicReference<>(Transmitter.capture());
        }

        @Override
        public void run() {
            SnapShot capture = captureRef.get();
            SnapShot backUp = Transmitter.replay(capture);
            try {
                runnable.run();
            } finally {
                Transmitter.restore(backUp);
            }
        }

        /**
         * 调用这个函数进行包装
         */
        public static EsRunnable getRunnable(Runnable runnable) {
            return new EsRunnable(runnable);
        }
    }
}

class Test {


    private static TransmittableThreadLocal<String> context = new TransmittableThreadLocal<>();
    private static ExecutorService pool = TtlExecutors.getTtlExecutorService(Executors.newFixedThreadPool(5));


    public static void main(String[] args)throws Exception {
        context.set("ershi");

        for (int i = 1; i <=10; i++) {
            if (i==2)
                context.set("yhd");
            if (i==7)
                context.set("ershi");
            pool.execute(new EsThreadLocal.EsRunnable(()->System.out.println(Thread.currentThread().getName()+" : " + context.get() )));
            TimeUnit.SECONDS.sleep(1);
        }
        pool.shutdown();
    }
}
